from rdkit import Chem, DataStructs
from rdkit.Chem import AllChem, rdChemReactions


def canonicalize_smiles(smiles):
    mol = Chem.MolFromSmiles(Chem.CanonSmiles(smiles))
    if mol is not None:
        return Chem.MolToSmiles(mol, isomericSmiles=False)
    else:
        return ""


def get_fingerprint(chemical_smiles, fingerprint_type):
    """input chemical smiles and fingerprint type to get its fingerprint;
    fingerprint type should be "Morgan2","Morgan3", "Morgan2withF" or "Morgan3withF\""""
    mol = Chem.MolFromSmiles(chemical_smiles)
    if fingerprint_type == "Morgan2":
        fp = AllChem.GetMorganFingerprint(mol, 2)
    elif fingerprint_type == "Morgan3":
        fp = AllChem.GetMorganFingerprint(mol, 3)
    elif fingerprint_type == "Morgan2withF":
        fp = AllChem.GetMorganFingerprint(mol, 2, useFeatures=True)
    elif fingerprint_type == "Morgan3withF":
        fp = AllChem.GetMorganFingerprint(mol, 3, useFeatures=True)
    return fp


def calculate_similarity(fp1, fp2, similarity_type):
    """input two chemical fingerprints and get their similarity"""
    if similarity_type == "Tanimoto":
        similarity = DataStructs.TanimotoSimilarity(fp1, fp2)
    elif similarity_type == "Dice":
        similarity = DataStructs.DiceSimilarity(fp1, fp2)
    elif similarity_type == "Tversky1":
        similarity = DataStructs.cDataStructs.TverskySimilarity(fp1, fp2, 1, 0.5)
    elif similarity_type == "Tversky2":
        similarity = DataStructs.cDataStructs.TverskySimilarity(fp1, fp2, 0.5, 1)
    return similarity


def apply_templates(reactant, template):
    """input chemical smiles and reaction templates to get possible products in a list"""
    try:
        rxn = rdChemReactions.ReactionFromSmarts(template)
        reacts = Chem.MolFromSmiles(reactant)
        product = rxn.RunReactant(reacts, 0)
        product_smiles_list = []
        for _ in range(len(product)):
            product_smiles = Chem.MolToSmiles(product[_][0], canonical=False, isomericSmiles=False)
            product_smiles_list.append(product_smiles)
        return product_smiles_list
    except (Exception,):
        return ""


# for precision calculation
def get_precision(top_n, prediction, answer):
    """input top_n range, prediction list and answer list to get top_n precison"""
    try:
        n = int(top_n)
        predict_n = prediction[:n]
        top_n_correct = [x for x in predict_n if x in answer]
        top_n_precision = len(top_n_correct) / len(predict_n)
        return top_n_precision
    except (Exception,):
        return ""


# for recall calculation
def get_recall(top_n, prediction, answer):
    """input top_n range, prediction list and answer list to get top_n recall"""
    try:
        n = int(top_n)
        predict_n = prediction[:n]
        top_n_correct = [x for x in predict_n if x in answer]
        top_n_recall = len(top_n_correct) / len(answer)
        return top_n_recall
    except (Exception,):
        return ""
